from ultralytics import YOLO
import os
import csv
import uproot
import numpy as np
from tqdm import tqdm

el_name = "H_Z1A1"
#el_name = "He_Z2A4"
#el_name = "C_Z6A12"
#el_name = "O_Z8A16"
#el_name = "D_Z1A2"
#el_name = "test_Plots"

# 关闭YOLO model输出
os.environ['YOLO_VERBOSE'] = 'False'

# 加载训练好的模型
model = YOLO("/home/mart/WorkSpace/Marvel/YOLO/train15/weights/best.pt")

def migdalYOLO(el_name):
    # 指定测试文件夹的路径
    folder_path = f"/home/mart/WorkSpace/Marvel/YOLO/NuclearRecoil/{el_name}"
    
    figure_list = os.listdir(folder_path)
    num_of_figures = len(figure_list)
    entry_num   = np.zeros(num_of_figures, dtype=int)
    er_num     = np.zeros(num_of_figures, dtype=int)
    nr_num     = np.zeros(num_of_figures, dtype=int)
    cosmic_num = np.zeros(num_of_figures, dtype=int)
    for ie, file_name in tqdm(enumerate(figure_list), total=num_of_figures):
    #for ie, file_name in enumerate(figure_list):
        file_path = os.path.join(folder_path, file_name)
        results = model(file_path)
        number_NR = 0
        number_ER = 0
        number_Cosmic = 0
        for r in results:
            boxes_cls_array = r.boxes.cls.cpu().numpy()
            #boxes_xywh_array = r.boxes.xywh.cpu().numpy()
            for i in range(len(boxes_cls_array)):
                #print(boxes_cls_array[i])
                if i == 0:
                    if boxes_cls_array[i] == 0:
                        number_ER += 1
                    elif boxes_cls_array[i] == 1:
                        number_NR += 1
                    elif boxes_cls_array[i] == 2:
                        number_Cosmic += 1
                else:
                    if boxes_cls_array[i] == 0:
                        number_ER += 1
                    elif boxes_cls_array[i] == 1:
                        number_NR += 1
                    elif boxes_cls_array[i] == 2:
                        number_Cosmic += 1
        entry_num[ie] = int(file_name.replace(".png",""))
        er_num[ie] = number_ER
        nr_num[ie] = number_NR
        cosmic_num[ie] = number_Cosmic
        #print(f"For entry {entry_num[ie]}: \tNR: {nr_num[ie]}, \tER:{er_num[ie]}, \tCosmic: {cosmic_num[ie]}")
    
    # 按照entry编号重排序
    seq_index = np.argsort(entry_num)
    entry_num = entry_num[seq_index]
    er_num = er_num[seq_index]
    nr_num = nr_num[seq_index]
    cosmic_num = cosmic_num[seq_index]
    
    # 将结果保存至ROOT文件
    with uproot.recreate(f'{folder_path}/../{el_name}_wide_yolo.root') as file:
        file["yolo_output"] = {
            "EntryNum": entry_num,
            "NRNum": nr_num,
            "ERNum": er_num,
            "CosmicNum": cosmic_num
        }


if __name__ == "__main__":
    el_list = ["H_Z1A1", "He_Z2A4", "C_Z6A12"]
    #el_list = ["H_Z1A1", "He_Z2A4", "C_Z6A12", "O_Z8A16"]
    #for el in el_list:
    migdalYOLO("CheckTails")
